"""
@Time: 2020/12/24 下午 1:21
@Author: jinzhuan
@File: bert_event.py
@Desc: 
"""

from cognlp import *
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import RandomSampler
from cognlp.core.metrics import EventMetric
from cognlp.core.trainer import Trainer
from cognlp.core.dataset import ACE2005Dataset
from cognlp.models.ee.bert_event import Bert4Event
from cognlp.utils.util import seed_everything

seed_everything()

torch.cuda.set_device(0)
device = torch.device('cuda')

trigger_vocabulary = Vocabulary.load('../../../cognlp/data/ee/ace2005/vocabulary.txt')
argument_vocabulary = Vocabulary.load('../../../cognlp/data/ee/ace2005/argument_vocabulary.txt')

train_dataset = ACE2005Dataset('../../../cognlp/data/ee/ace2005/data/train.json', trigger_vocabulary,
                               argument_vocabulary)
train_samples_weight = train_dataset.get_samples_weight()
train_sampler = torch.utils.data.WeightedRandomSampler(train_samples_weight, len(train_samples_weight))

dev_dataset = ACE2005Dataset('../../../cognlp/data/ee/ace2005/data/dev.json', trigger_vocabulary, argument_vocabulary)

test_dataset = ACE2005Dataset('../../../cognlp/data/ee/ace2005/data/test.json', trigger_vocabulary, argument_vocabulary)

model = Bert4Event(trigger_vocabulary=trigger_vocabulary, argument_vocabulary=argument_vocabulary)
loss = nn.CrossEntropyLoss(ignore_index=0)
optimizer = optim.Adam(model.parameters(), lr=0.00005)
metric = EventMetric(trigger_vocabulary, argument_vocabulary)

trainer = Trainer(model,
                  train_dataset,
                  dev_data=test_dataset,
                  n_epochs=30,
                  batch_size=20,
                  loss=loss,
                  optimizer=optimizer,
                  scheduler=None,
                  metrics=metric,
                  train_sampler=train_sampler,
                  dev_sampler=None,
                  drop_last=False,
                  gradient_accumulation_steps=1,
                  num_workers=None,
                  save_path=None,
                  save_file=None,
                  print_every=None,
                  scheduler_steps=None,
                  validate_steps=None,
                  save_steps=None,
                  grad_norm=1.0,
                  use_tqdm=True,
                  device=device,
                  device_ids=[0],
                  collate_fn=train_dataset.to_dict,
                  callbacks=None,
                  metric_key=None,
                  writer_path='../../../cognlp/data/ee/ace2005/tensorboard',
                  fp16=False,
                  fp16_opt_level='O1',
                  checkpoint_path=None,
                  task='ace2005-event-test',
                  logger_path='../../../cognlp/data/ee/ace2005/logger')

trainer.train()
# tester = Tester(model,
#                 model_path='../../../cognlp/data/ee/ace2005/model/ace2005-event/2020-12-17-18:14:48/checkpoint-820/model.pt',
#                 batch_size=32, sampler=None,
#                 drop_last=False, num_workers=0, print_every=1000,
#                 dev_data=dev_dataset, metrics=metric, metric_key=None, use_tqdm=True, device=device,
#                 callbacks=None, check_code_level=0, device_ids=[4, 5])
# tester.test()
print(1)
